from sklearn.metrics import accuracy_score, f1_score, confusion_matrix
import seaborn as sns
import matplotlib.pyplot as plt
import pandas as pd
from sklearn.model_selection import train_test_split
import torch
from torch.utils.data import Dataset, DataLoader
from transformers import BertTokenizer, BertForSequenceClassification
from torch.optim import AdamW
from torch.nn.utils.rnn import pad_sequence
from torch.utils.tensorboard import SummaryWriter


# 读取CSV文件
df = pd.read_csv('data/cmu_mosei/metadata/merged_data_clean_small.csv')
texts = df['english'].fillna('').astype(str).tolist()
labels = df['label'].astype(int).tolist()
print(len(labels))

# 加载BERT tokenizer
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
# 划分训练集和验证集
X_train, X_val, y_train, y_val = train_test_split(texts, labels, test_size=0.2, stratify=labels, random_state=42)


# 自定义数据集类
class TextDataset(Dataset):
    def __init__(self, texts, labels, tokenizer):
        self.texts = texts
        self.labels = labels
        self.tokenizer = tokenizer

    def __len__(self):
        return len(self.texts)

    def __getitem__(self, idx):
        text = self.texts[idx]
        label = self.labels[idx]
        encoding = self.tokenizer(text, padding=False, truncation=True, max_length=1024, return_tensors='pt')
        return {
            'input_ids': encoding['input_ids'].squeeze(),
            'attention_mask': encoding['attention_mask'].squeeze(),
            'labels': torch.tensor(label, dtype=torch.long)
        }


# 自定义collate函数
def collate_fn(batch):
    input_ids = [item['input_ids'] for item in batch]
    attention_mask = [item['attention_mask'] for item in batch]
    labels = [item['labels'] for item in batch]
    return {
        'input_ids': pad_sequence(input_ids, batch_first=True, padding_value=0),
        'attention_mask': pad_sequence(attention_mask, batch_first=True, padding_value=0),
        'labels': torch.tensor(labels, dtype=torch.long),
    }


# 训练和验证函数
def train_one_epoch(model, dataloader, optimizer, device, writer, epoch, current_learning_rate):
    model.train()
    running_loss = 0.0
    for step, batch in enumerate(dataloader):
        optimizer.zero_grad()
        input_ids = batch['input_ids'].to(device)
        attention_mask = batch['attention_mask'].to(device)
        labels = batch['labels'].to(device)

        outputs = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels)
        loss = outputs.loss
        loss.backward()
        optimizer.step()
        running_loss += loss.item()

        # 记录 Step Loss 和 Learning Rate 到 TensorBoard
        writer.add_scalar('Train/Step_Loss', loss.item(), epoch * len(dataloader) + step)
        writer.add_scalar('Train/Learning_Rate', current_learning_rate, epoch * len(dataloader) + step)

    # 计算并返回 Epoch Loss
    epoch_loss = running_loss / len(dataloader)
    writer.add_scalar('Train/Epoch_Loss', epoch_loss, epoch)
    return epoch_loss


def validate_one_epoch(model, dataloader, device, writer, epoch):
    model.eval()
    running_loss = 0.0
    with torch.no_grad():
        for step, batch in enumerate(dataloader):
            input_ids = batch['input_ids'].to(device)
            attention_mask = batch['attention_mask'].to(device)
            labels = batch['labels'].to(device)

            outputs = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels)
            loss = outputs.loss
            running_loss += loss.item()

            # 记录 Step Loss 到 TensorBoard
            writer.add_scalar('Validation/Step_Loss', loss.item(), epoch * len(dataloader) + step)

    # 计算并返回 Epoch Loss
    epoch_loss = running_loss / len(dataloader)
    writer.add_scalar('Validation/Epoch_Loss', epoch_loss, epoch)
    return epoch_loss


# 训练主函数
def train_and_evaluate(epochs, patience, learning_rate, experiment_index):
    # 加载BERT预训练模型
    model = BertForSequenceClassification.from_pretrained('bert-base-uncased', num_labels=7)
    train_dataset = TextDataset(X_train, y_train, tokenizer)
    val_dataset = TextDataset(X_val, y_val, tokenizer)

    train_dataloader = DataLoader(train_dataset, batch_size=32, shuffle=True, collate_fn=collate_fn)
    val_dataloader = DataLoader(val_dataset, batch_size=32, collate_fn=collate_fn)

    # 设置设备
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)

    # 优化器
    optimizer = AdamW(model.parameters(), lr=learning_rate, weight_decay=0.01)
    best_val_loss = float('inf')
    counter = 0

    # 创建独立的TensorBoard日志，实验名包含学习率
    writer = SummaryWriter(log_dir=f'./logs/mer_lr_{learning_rate:.2e}')

    for epoch in range(epochs):
        train_loss = train_one_epoch(model, train_dataloader, optimizer, device, writer, epoch, learning_rate)
        val_loss = validate_one_epoch(model, val_dataloader, device, writer, epoch)

        print(
            f"Experiment {experiment_index} | Epoch {epoch + 1}/{epochs}, Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f}")

        # 早停逻辑
        if val_loss < best_val_loss:
            best_val_loss = val_loss
            counter = 0
            # 保存模型权重
            model.save_pretrained(f'./bert/mer_lr_{learning_rate:.2e}')
            tokenizer.save_pretrained(f'./bert/mer_lr_{learning_rate:.2e}')
        else:
            counter += 1
            if counter >= patience:
                print("Early stopping triggered.")
                break

    writer.close()
    return best_val_loss


# 预测函数
def predict_batch(model, texts, batch_size=16):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)
    model.eval()

    predictions = []
    for i in range(0, len(texts), batch_size):
        batch_texts = texts[i:i + batch_size]
        inputs = tokenizer(batch_texts, return_tensors='pt', padding=True, truncation=True, max_length=512)
        inputs = {key: value.to(device) for key, value in inputs.items()}

        with torch.no_grad():
            outputs = model(**inputs)
            logits = outputs.logits
            predicted_labels = torch.argmax(logits, dim=-1)
            predictions.extend(predicted_labels.cpu().numpy())

    return predictions


# 新增的预测调用封装函数
def evaluate_from_csv(csv_path, model_path='./bert/experiment_mosei_15_lr_1e-05'):
    # 读取CSV文件
    df = pd.read_csv(csv_path)
    texts = df['english'].fillna('').astype(str).tolist()
    true_labels = df['label'].astype(int).tolist()

    # 加载BERT模型
    model = BertForSequenceClassification.from_pretrained(model_path)

    # 批量预测
    predicted_labels = predict_batch(model, texts)

    # 计算评估指标
    accuracy = accuracy_score(true_labels, predicted_labels)
    f1 = f1_score(true_labels, predicted_labels, average='weighted')
    f1_macro = f1_score(true_labels, predicted_labels, average='macro')
    cm = confusion_matrix(true_labels, predicted_labels)

    # 输出评估结果
    print(f"Accuracy: {accuracy:.4f}")
    print(f"Weighted F1 Score: {f1:.4f}")
    print(f"Macro F1 Score: {f1_macro:.4f}")

    # 绘制混淆矩阵
    class_names = [str(i) for i in range(7)]  # 假设标签是 0, 1, 2, ..., 5
    plot_confusion_matrix(cm, class_names)


# 混淆矩阵绘制函数
def plot_confusion_matrix(cm, class_names):
    plt.figure(figsize=(8, 6))
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', xticklabels=class_names, yticklabels=class_names)
    plt.xlabel('Predicted')
    plt.ylabel('True')
    plt.title('Confusion Matrix')
    plt.show()


# 主程序
if __name__ == '__main__':
    epochs = 100
    patience = 50
    learning_rates = [1e-7,2e-7,3e-7,5e-7]
    experiment_results = []

    for i in range(len(learning_rates)):
        # 调整学习率
        current_learning_rate = learning_rates[i]

        print(f"\nStarting Experiment {i + 1} with Learning Rate {current_learning_rate:}")
        best_val_loss = train_and_evaluate(epochs, patience, current_learning_rate, i + 1)
        experiment_results.append((i + 1, current_learning_rate, best_val_loss))


    print("Training Complete! Evaluating the best model:")
    best_model_path = './bert/small_mosei_lr_1.00e-07'
    print("Evaluating the max 1e-4 model:")
    # csv_path = 'data/mer/text/merged_file_with_labels.csv'
    csv_path = 'data/cmu_mosei/metadata/merged_data_clean_small.csv'
    evaluate_from_csv(csv_path, model_path=best_model_path)
    print("Evaluating the max 2e-4 model:")
    evaluate_from_csv(csv_path, model_path='./bert/small_mosei_lr_2.00e-07')
    print("Evaluating the max 3e-4 model:")
    evaluate_from_csv(csv_path, model_path='./bert/small_mosei_lr_3.00e-07')
    print("Evaluating the max 5e-4 model:")
    evaluate_from_csv(csv_path, model_path='./bert/small_mosei_lr_5.00e-07')
